# @title 绘图工具

# 获取Noto JP字体以显示日文字符
# !sudo apt update
# !apt-get install fonts-noto-cjk  # 用于Noto Sans CJK JP字体
# !apt-get install fonts-source-han-sans-jp # 用于Source Han Sans（日文）

import json
import random
import io
import ast
from PIL import Image, ImageDraw, ImageFont
from PIL import ImageColor
from IPython.display import Markdown, display
from openai import OpenAI
import os
import base64


def plot_text_bounding_boxes(image_path, bounding_boxes):
    """
    在图像上绘制边界框，并为每个边界框标记名称，使用PIL库、标准化坐标和不同颜色。

    参数:
        image_path: 图像文件的路径。
        bounding_boxes: 包含对象名称及其在标准化[y1 x1 y2 x2]格式中位置的边界框列表。
    """

    # 加载图像
    img = Image.open(image_path)
    width, height = img.size
    print(img.size)
    # 创建绘图对象
    draw = ImageDraw.Draw(img)

    # 解析markdown围栏
    bounding_boxes = parse_json(bounding_boxes)

    # 注意：Windows系统上可能需要修改字体路径
    try:
        font = ImageFont.truetype("NotoSansCJK-Regular.ttc", size=10)
    except IOError:
        # 回退到默认字体
        font = ImageFont.load_default()

    # 遍历边界框
    for i, bounding_box in enumerate(ast.literal_eval(bounding_boxes)):
      color = 'green'

      # 将标准化坐标转换为绝对坐标
      abs_y1 = int(bounding_box["bbox_2d"][1]/999 * height)
      abs_x1 = int(bounding_box["bbox_2d"][0]/999 * width)
      abs_y2 = int(bounding_box["bbox_2d"][3]/999 * height)
      abs_x2 = int(bounding_box["bbox_2d"][2]/999 * width)

      if abs_x1 > abs_x2:
        abs_x1, abs_x2 = abs_x2, abs_x1

      if abs_y1 > abs_y2:
        abs_y1, abs_y2 = abs_y2, abs_y1

      # 绘制边界框
      draw.rectangle(
          ((abs_x1, abs_y1), (abs_x2, abs_y2)), outline=color, width=1
      )

      # 绘制文本
      if "text_content" in bounding_box:
        draw.text((abs_x1, abs_y2), bounding_box["text_content"], fill=color, font=font)

    # 显示图像
    img.show()

# @title 解析JSON输出
def parse_json(json_output):
    # 解析markdown围栏
    lines = json_output.splitlines()
    for i, line in enumerate(lines):
        if line == "```json":
            json_output = "\n".join(lines[i+1:])  # 移除"```json"之前的所有内容
            json_output = json_output.split("```")[0]  # 移除结束"```"之后的所有内容
            break  # 一旦找到"```json"就退出循环
    return json_output

# @title 推理函数（本地版）
def inference(image_path, prompt, sys_prompt="You are a helpful assistant.", max_new_tokens=4096, return_input=False):
    image = Image.open(image_path)
    image_local_path = "file://" + image_path
    messages = [
        # 跳过系统提示
        # {"role": "system", "content": sys_prompt},
        {"role": "user", "content": [
                {"type": "text", "text": prompt},
                {"image": image_local_path},
            ]
        },
    ]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    print("text:", text)
    # image_inputs, video_inputs = process_vision_info([messages])
    inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
    inputs = inputs.to('cuda')

    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens)
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
    output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
    if return_input:
        return output_text[0], inputs
    else:
        return output_text[0]
    



# base64编码格式
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


# @title API推理函数
def inference_with_api(image_path, prompt, sys_prompt="你是一个有帮助的助手。", model_id="qwen3-vl-235b-a22b-instruct", min_pixels=512*32*32, max_pixels=2048*32*32):
    base64_image = encode_image(image_path)
    client = OpenAI(
        # 如果未配置环境变量，请用Dashscope API密钥替换下行：api_key="sk-xxx"。访问 https://bailian.console.alibabacloud.com/?apiKey=1 获取
        api_key="sk-3fe722ea8cd249beb3bf3da52099484b",
        base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
    )


    messages=[
        # 跳过系统提示
        # {
        #     "role": "system",
        #     "content": [{"type":"text","text": sys_prompt}]},
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "min_pixels": min_pixels,
                    "max_pixels": max_pixels,
                    # 传入BASE64图像数据。注意图像格式（即image/{format}）必须与支持的图像列表中的Content Type匹配。"f"是字符串格式化方法。
                    # PNG图像:  f"data:image/png;base64,{base64_image}"
                    # JPEG图像: f"data:image/jpeg;base64,{base64_image}"
                    # WEBP图像: f"data:image/webp;base64,{base64_image}"
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                },
                {"type": "text", "text": prompt},
            ],
        }
    ]
    completion = client.chat.completions.create(
        model = model_id,
        messages = messages,
       
    )
    return completion.choices[0].message.content